﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Microsoft.Media.ISO
{
    public static class NALUReader
    {
        public static IList<NALU> ReadAnnexB([System.Runtime.InteropServices.WindowsRuntime.ReadOnlyArray] byte[] bytes)
        {
            var result = new List<NALU>();

            NALU nalu = null;
            int i = 0;
            int offset = i;
            do
            {
                if ((i >= bytes.Length) || (i > 1 && bytes[i] == 0x01 && bytes[i - 1] == 0x00 && bytes[i - 2] == 0x00))
                {
                    if (nalu != null) // save the last nalu
                    {
                        int length = i - offset - 3;
                        if (i > 2 && bytes[i - 3] == 0x00) // the start code might be 4 bytes. If 4th byte ago was zero, assume so.
                        {
                            length--;
                        }
                        nalu.Data = new byte[length];
                        System.Buffer.BlockCopy(bytes, offset, nalu.Data, 0, length);
                        result.Add(nalu);
                    }
                    if (i >= bytes.Length) break;

                    nalu = new NALU();
                    var b = bytes[i++];
                    nalu.IsReferenced = ((b >> 5) & 0x03) == 1; // bits 1-2 indicate weather is NALU is referenced by another NALU
                    nalu.Type = (NALUType)(b & 0x1f); // bits 3 through 7 contains the NALU type
                    offset = i + 1;
                }
                i++;
            } while (true);

            return result;
        }

        public static IList<NALU> ReadAvvC([System.Runtime.InteropServices.WindowsRuntime.ReadOnlyArray] byte[] bytes, int naluSizeLength)
        {
            var result = new List<NALU>();

            int i = 0;
            while (i < bytes.Length)
            {
                int size = 0;
                for (int j = 0; j < naluSizeLength; j++)
                {
                    size = (size << 8) + bytes[i + j];
                }
                i += naluSizeLength;
                
                var nalu = new NALU();
                var b = bytes[i];
                nalu.IsReferenced = ((b >> 5) & 0x03) == 1; // bits 1-2 indicate weather is NALU is referenced by another NALU
                nalu.Type = (NALUType)(b & 0x1f); // bits 3 through 7 contains the NALU type

                nalu.Data = new byte[size];
                System.Buffer.BlockCopy(bytes, i, nalu.Data, 0, size);
                
                result.Add(nalu);

                i += size;
            }

            return result;
        }
    }

    public sealed class NALU
    {
        public bool IsVideoCodingLayer
        {
            get { return Type <= NALUType.CodedSliceOfAnIDRPicture; }
        }

        public bool IsReferenced { get; set; }

        public NALUType Type { get; set; }

        public byte[] Data { get; set; }
    }

    public enum NALUType : uint
    {
        Unspecified = 0,
        CodedSliceOfANonIDRPicture = 1,
        CodedSliceDataPartitionA = 2,
        CodedSliceDataPartitionB = 3,
        CodedSliceDataPartitionC = 4,
        CodedSliceOfAnIDRPicture = 5,
        SupplementalEnhancementInformation = 6,
        SequenceParameterSet = 7,
        PictureParameterSet = 8,
        AccessUnitDelimiter = 9,
        EndOfSequence = 10,
        EndOfStream = 11,
        FillerData = 12,
        SequenceParameterSetExtension = 13,
        PrefixNALUnit = 14,
        SubsetSequenceParameterSet = 15,
        Reserved1 = 16,
        Reserved2 = 17,
        Reserved3 = 18,
        CodedSliceOfAnAuxiliaryCodedPictureWithoutPartitioning = 19,
        CodedSliceExtension = 20,
        CodedSliceExtensionForDepthViewComponents = 21,
        Reserved4 = 22,
        Reserved5 = 23,
        Unspecified1 = 24,
        Unspecified2 = 25,
        Unspecified3 = 26,
        Unspecified4 = 27,
        Unspecified5 = 28,
        Unspecified6 = 29,
        Unspecified7 = 30,
        Unspecified8 = 31
    }
}
